{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "50bd7e7c",
   "metadata": {},
   "source": [
    "# Unit Test Writer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad339d11",
   "metadata": {},
   "source": [
    "### Welcome to the Unit Test Writer - an AI-powered tool that automatically generates comprehensive unit tests for your code across multiple programming languages. Simply paste your code, select your language, and let AI create thorough test suites in seconds!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27e4e719",
   "metadata": {},
   "outputs": [],
   "source": [
    "# imports\n",
    "\n",
    "import os\n",
    "import io\n",
    "import sys\n",
    "from dotenv import load_dotenv\n",
    "from openai import OpenAI\n",
    "import gradio as gr\n",
    "import subprocess\n",
    "from IPython.display import Markdown, display"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3572f5fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "load_dotenv(override=True)\n",
    "openai_api_key = os.getenv('OPENAI_API_KEY')\n",
    "anthropic_api_key = os.getenv('ANTHROPIC_API_KEY')\n",
    "google_api_key = os.getenv('GOOGLE_API_KEY')\n",
    "grok_api_key = os.getenv('GROK_API_KEY')\n",
    "groq_api_key = os.getenv('GROQ_API_KEY')\n",
    "openrouter_api_key = os.getenv('OPENROUTER_API_KEY')\n",
    "\n",
    "if openai_api_key:\n",
    "    print(f\"OpenAI API Key exists and begins {openai_api_key[:8]}\")\n",
    "else:\n",
    "    print(\"OpenAI API Key not set\")\n",
    "    \n",
    "if anthropic_api_key:\n",
    "    print(f\"Anthropic API Key exists and begins {anthropic_api_key[:7]}\")\n",
    "else:\n",
    "    print(\"Anthropic API Key not set (and this is optional)\")\n",
    "\n",
    "if google_api_key:\n",
    "    print(f\"Google API Key exists and begins {google_api_key[:2]}\")\n",
    "else:\n",
    "    print(\"Google API Key not set (and this is optional)\")\n",
    "\n",
    "if grok_api_key:\n",
    "    print(f\"Grok API Key exists and begins {grok_api_key[:4]}\")\n",
    "else:\n",
    "    print(\"Grok API Key not set (and this is optional)\")\n",
    "\n",
    "if groq_api_key:\n",
    "    print(f\"Groq API Key exists and begins {groq_api_key[:4]}\")\n",
    "else:\n",
    "    print(\"Groq API Key not set (and this is optional)\")\n",
    "\n",
    "if openrouter_api_key:\n",
    "    print(f\"OpenRouter API Key exists and begins {openrouter_api_key[:6]}\")\n",
    "else:\n",
    "    print(\"OpenRouter API Key not set (and this is optional)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05293821",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Connect to client libraries\n",
    "\n",
    "openai = OpenAI()\n",
    "\n",
    "anthropic_url = \"https://api.anthropic.com/v1/\"\n",
    "gemini_url = \"https://generativelanguage.googleapis.com/v1beta/openai/\"\n",
    "grok_url = \"https://api.x.ai/v1\"\n",
    "groq_url = \"https://api.groq.com/openai/v1\"\n",
    "ollama_url = \"http://localhost:11434/v1\"\n",
    "openrouter_url = \"https://openrouter.ai/api/v1\"\n",
    "\n",
    "anthropic = OpenAI(api_key=anthropic_api_key, base_url=anthropic_url)\n",
    "gemini = OpenAI(api_key=google_api_key, base_url=gemini_url)\n",
    "grok = OpenAI(api_key=grok_api_key, base_url=grok_url)\n",
    "groq = OpenAI(api_key=groq_api_key, base_url=groq_url)\n",
    "ollama = OpenAI(api_key=\"ollama\", base_url=ollama_url)\n",
    "openrouter = OpenAI(api_key=openrouter_api_key, base_url=openrouter_url)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41eec7d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "models = [\"gpt-5\", \"claude-sonnet-4-5-20250929\", \"grok-4\", \"gemini-2.5-pro\", \"qwen2.5-coder\", \"deepseek-coder-v2\", \"gpt-oss:20b\", \"qwen/qwen3-coder-30b-a3b-instruct\", \"openai/gpt-oss-120b\", ]\n",
    "\n",
    "clients = {\"gpt-5\": openai, \"claude-sonnet-4-5-20250929\": anthropic, \"grok-4\": grok, \"gemini-2.5-pro\": gemini, \"openai/gpt-oss-120b\": groq, \"qwen2.5-coder\": ollama, \"deepseek-coder-v2\": ollama, \"gpt-oss:20b\": ollama, \"qwen/qwen3-coder-30b-a3b-instruct\": openrouter}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05ca8fcc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to generate unit tests\n",
    "def generate_unit_tests(code, language, model_name):\n",
    "    \"\"\"Generate unit tests for the provided code\"\"\"\n",
    "    \n",
    "    if not code.strip():\n",
    "        return \"Please provide some code to generate tests for.\"\n",
    "    \n",
    "    # Select the appropriate client\n",
    "    client = clients.get(model_name, openai)\n",
    "    \n",
    "    # Create the prompt based on language\n",
    "    test_frameworks = {\n",
    "        \"Python\": \"pytest\",\n",
    "        \"JavaScript\": \"Jest\",\n",
    "        \"TypeScript\": \"Jest\",\n",
    "        \"Java\": \"JUnit\",\n",
    "        \"C#\": \"NUnit\",\n",
    "        \"Go\": \"testing package\",\n",
    "        \"Ruby\": \"RSpec\",\n",
    "        \"PHP\": \"PHPUnit\",\n",
    "        \"Rust\": \"built-in test framework\"\n",
    "    }\n",
    "    \n",
    "    framework = test_frameworks.get(language, \"appropriate testing framework\")\n",
    "    \n",
    "    prompt = f\"\"\"You are a unit test expert. Generate comprehensive unit tests for the following {language} code using {framework}.\n",
    "\n",
    "                    Code to test:\n",
    "                    ```{language.lower()}\n",
    "                    {code}\n",
    "                    ```\n",
    "\n",
    "                    Requirements:\n",
    "                    - Create thorough unit tests covering normal cases, edge cases, and error cases\n",
    "                    - Use {framework} syntax and best practices\n",
    "                    - Include clear test names that describe what is being tested\n",
    "                    - Add comments explaining complex test scenarios\n",
    "                    - Ensure tests are independent and can run in any order\n",
    "\n",
    "                    Provide ONLY the test code, no explanations.\"\"\"\n",
    "\n",
    "    try:\n",
    "        # Call the API\n",
    "        response = client.chat.completions.create(\n",
    "            model=model_name,\n",
    "            messages=[\n",
    "                {\"role\": \"user\", \"content\": prompt}\n",
    "            ],\n",
    "            temperature=0.7,\n",
    "        )\n",
    "        \n",
    "        return response.choices[0].message.content\n",
    "        \n",
    "    except Exception as e:\n",
    "        return f\"Error generating tests: {str(e)}\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d51e13d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create Gradio interface\n",
    "def create_interface():\n",
    "    with gr.Blocks(title=\"Unit Test Generator\") as demo:\n",
    "        gr.Markdown(\"# 🧪 Unit Test Generator\")\n",
    "        gr.Markdown(\"Paste your code, select the language, and generate comprehensive unit tests!\")\n",
    "        \n",
    "        with gr.Row():\n",
    "            with gr.Column():\n",
    "                code_input = gr.TextArea(\n",
    "                    label=\"Your Code\",\n",
    "                    placeholder=\"Paste your code here...\",\n",
    "                    lines=15\n",
    "                )\n",
    "                \n",
    "                with gr.Row():\n",
    "                    language_dropdown = gr.Dropdown(\n",
    "                        choices=[\"Python\", \"JavaScript\", \"TypeScript\", \"Java\", \"C#\", \"Go\", \"Ruby\", \"PHP\", \"Rust\"],\n",
    "                        value=\"Python\",\n",
    "                        label=\"Language\"\n",
    "                    )\n",
    "                    \n",
    "                    model_dropdown = gr.Dropdown(\n",
    "                        choices=models,\n",
    "                        value=\"gpt-5\",\n",
    "                        label=\"Model\"\n",
    "                    )\n",
    "                \n",
    "                generate_btn = gr.Button(\"Generate Unit Tests\", variant=\"primary\")\n",
    "            \n",
    "            with gr.Column():\n",
    "                output = gr.TextArea(\n",
    "                    label=\"Generated Unit Tests\",\n",
    "                    lines=15\n",
    "                )\n",
    "        \n",
    "        # Example\n",
    "        gr.Markdown(\"### Example\")\n",
    "        gr.Examples(\n",
    "            examples=[\n",
    "                [\"\"\"def add(a, b):\n",
    "    return a + b\n",
    "\n",
    "def multiply(a, b):\n",
    "    return a * b\n",
    "\n",
    "def divide(a, b):\n",
    "    if b == 0:\n",
    "        raise ValueError(\"Cannot divide by zero\")\n",
    "    return a / b\"\"\", \"Python\", \"gpt-5\"],\n",
    "                [\"\"\"function isPalindrome(str) {\n",
    "    const cleaned = str.toLowerCase().replace(/[^a-z0-9]/g, '');\n",
    "    return cleaned === cleaned.split('').reverse().join('');\n",
    "}\"\"\", \"JavaScript\", \"gpt-5\"]\n",
    "            ],\n",
    "            inputs=[code_input, language_dropdown, model_dropdown]\n",
    "        )\n",
    "        \n",
    "        # Connect the button\n",
    "        generate_btn.click(\n",
    "            fn=generate_unit_tests,\n",
    "            inputs=[code_input, language_dropdown, model_dropdown],\n",
    "            outputs=output\n",
    "        )\n",
    "    \n",
    "    return demo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8eff8ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Launch the interface\n",
    "demo = create_interface()\n",
    "demo.launch(share=False)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
